from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
from datasets import Dataset, Audio
import torch
import logging

# 配置日志
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

def test_whisper():
    """测试Whisper的具体实现"""

    device = 'cuda:9' if torch.cuda.is_available() else 'cpu'
    logger.info(f"使用设备: {device}")

    # 创建pipeline - 完全按照参考实现
    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
    model_id = 'openai/whisper-large-v3'

    logger.info("加载模型...")
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        model_id,
        torch_dtype=torch_dtype,
        low_cpu_mem_usage=True,
        use_safetensors=True
    )
    model.to(device)

    processor = AutoProcessor.from_pretrained(model_id)

    pipe = pipeline(
        'automatic-speech-recognition',
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=processor.feature_extractor,
        chunk_length_s=30,
        max_new_tokens=1024,
        batch_size=16,
        return_timestamps='word',
        torch_dtype=torch_dtype,
        device=device
    )

    kwargs = {'language': 'chinese', 'task': 'transcribe'}

    # 测试文件路径
    audio_file = "/root/code/drwork/ai2/uploads/test_audio.mp3"  # 假设有一个测试文件

    try:
        logger.info(f"开始转录文件: {audio_file}")

        # 创建数据集
        ds = Dataset.from_dict({'audio': [audio_file]}).cast_column(
            'audio', Audio(sampling_rate=16000))

        # 处理数据
        res = []
        for data in ds:
            sample = data['audio']
            logger.info(f"处理音频片段，时长: {sample['array'].size / sample['sampling_rate']:.2f}秒")

            # 进行转录
            pred = pipe(sample.copy(), generate_kwargs=kwargs)
            logger.info(f"转录结果类型: {type(pred)}")
            logger.info(f"转录结果键: {pred.keys() if isinstance(pred, dict) else 'Not a dict'}")

            if 'chunks' in pred:
                logger.info(f"Chunks类型: {type(pred['chunks'])}")
                logger.info(f"Chunks数量: {len(pred['chunks'])}")

                # 检查第一个chunk
                if pred['chunks']:
                    first_chunk = pred['chunks'][0]
                    logger.info(f"第一个chunk类型: {type(first_chunk)}")
                    logger.info(f"第一个chunk键: {first_chunk.keys() if isinstance(first_chunk, dict) else 'Not a dict'}")
                    logger.info(f"第一个chunk内容: {first_chunk}")

            res.append(pred['chunks'])

        logger.info(f"最终结果: {res}")
        return res

    except Exception as e:
        logger.error(f"错误发生: {e}")
        import traceback
        traceback.print_exc()
        raise e

if __name__ == "__main__":
    test_whisper()